package algos;

import java.io.IOException;
import java.util.*;
import utils.Pair;

import static utils.Common.*;

public class Engine {

    public static final int WORDS_TO_SUGGEST = 10;

    private final IVocabulary vocabulary;
    private final List<NGramModel> models;
    private final EstimationFunction estimation;

    private Engine(Vocabulary vocabulary, List<NGramModel> models, EstimationFunction estimation) {
        this.vocabulary = vocabulary;
        this.models = models;
        this.estimation = estimation;
    }

    public IVocabulary getVocabulary() {
        return vocabulary;
    }


    private void startPredict(List<String> words) {
        for (NGramModel model : models) {
            model.startPredict(words);
        }
    }

    private double predict(int word) {
        return estimation.p(word);
    }

    private void stopPredict() {
        for (NGramModel model : models) {

            model.stopPredict();
        }
    }



    public List<Pair<Double, String>> get(List<String> l, String prefix) {
        startPredict(l);
        List<Pair<Double, String>> answer = newArrayList(WORDS_TO_SUGGEST);
        for (int word : vocabulary.getWordsWithPrefix(prefix)) {
            double possib = predict(word);
            if (answer.size() == WORDS_TO_SUGGEST && answer.get(WORDS_TO_SUGGEST - 1).first > -possib) {
                answer.remove(WORDS_TO_SUGGEST - 1);
            }
            if (answer.size() < WORDS_TO_SUGGEST) {
                answer.add(makePair(-possib, vocabulary.getWordAt(word)));
                Collections.sort(answer);
            }
        }
        for (int i = 0, n = answer.size(); i < n; i++) {
            answer.set(i, makePair(-answer.get(i).first, answer.get(i).second));
        }
        stopPredict();
        return answer;
    }


    public static Engine create(List<Pair<String, Double>> ngramfiles) {
        Vocabulary vocabulary = new Vocabulary();
        for (Pair<String, Double> ngramfile : ngramfiles) {
            vocabulary.addWordsFromFile(ngramfile.first);
        }
        vocabulary.build();
        List<EstimationFunction> laplaces = newArrayList();
        List<Function<Integer, Double>> coefs = newArrayList();
        List<NGramModel> models = newArrayList();
        for (Pair<String, Double> ngramfile : ngramfiles) {
            NGramModel ngm = null;
            try {
                ngm = NGramModel.createFrom(ngramfile.first, vocabulary);
            } catch (IOException e) {
                System.err.println(ngramfile + ": " + e.getMessage());
                e.printStackTrace();
            }
            if (ngm == null) continue;
            models.add(ngm);
            laplaces.add(new LaplaceLaw(ngm));
            coefs.add(Functions.<Integer, Double>constantFunction(ngramfile.second));
        }
        EstimationFunction estimation = new GeneralLinearInterpolation(coefs, laplaces);
        return new Engine(vocabulary, models, estimation);
    }

}
